#include "RayTracer.h"

#include "KdTreeMemoryManager.h"
#include "TriangleListMemoryManager.h"
#include "Intersections.h"
#include "KdTraversalStack.h"
#include "LightSource.h"
#include "ShadowRay.h"
#include <assert.h>

extern BITMAP* buffer;
extern KdTreeMemoryManager kdtmm;
extern TriangleListMemoryManager tlmm;
extern SPoint viewPoint;
extern KdTraversalStack stack;

static float calculateCost(Triangle** triangles, const int numTriangles,
		const AABB& aabb, const float splitpos, const int axis, int& numLeft,
		int&numRight)
{
	const float leftSize = aabb.getLeftHullSize(splitpos, axis);
	const float rightSize = aabb.getRightHullSize(splitpos, axis);
	const float size = .5 / aabb.getHullSize();

	numLeft = numRight = 0;
	for (int i = 0; i < numTriangles; i++)
	{
		if (triangles[i]->getMin(axis) < splitpos)
			numLeft++;
		if (triangles[i]->getMax(axis) >= splitpos)
			numRight++;
	}

	return TRAVERSALCOST + INTERSECTCOST * (leftSize * size * numLeft
			+ rightSize * size * numRight);
}

static void _buildKdTree(KdNode* n, Triangle** tl, const int numTriangles,
		const AABB& aabb, int depth)
{
	// set the axis of the node
	const int axis = aabb.getMaxAxis();
	n->setAxis(axis);

	// find the split-plane and the number of triangles on each side
	int numLeft = 0;
	int numRight = 0;
	float split = MAXFLOAT;
	float bestcost = MAXFLOAT;
	int tmpNumLeft, tmpNumRight;
	for (int i = 0; i < numTriangles; i++)
	{
		const Triangle* triangle = tl[i];
		const float leftExtreme = triangle->getMin(axis);
		const float rightExtreme = triangle->getMax(axis);
		if (leftExtreme > aabb.getMin().get(axis))
		{
			float cost = calculateCost(tl, numTriangles, aabb, leftExtreme,
					axis, tmpNumLeft, tmpNumRight);
			if (cost < bestcost)
			{
				bestcost = cost;
				split = leftExtreme;
				numLeft = tmpNumLeft;
				numRight = tmpNumRight;
			}
		}

		if (rightExtreme < aabb.getMax().get(axis))
		{
			float cost = calculateCost(tl, numTriangles, aabb, rightExtreme,
					axis, tmpNumLeft, tmpNumRight);
			if (cost < bestcost)
			{
				bestcost = cost;
				split = rightExtreme;
				numLeft = tmpNumLeft;
				numRight = tmpNumRight;
			}
		}
	}

	// stop-criterions
	if (numTriangles <= MAXTRIANGLESPERKDNODE || depth >= MAXKDTREEDEPTH
			|| bestcost >= numTriangles * INTERSECTCOST)
	{
		// add all triangles to this node
		Triangle** t = new (tlmm.request((numTriangles + 1))) Triangle*[numTriangles + 1];
		for (int i = 0; i < numTriangles; i++)
			*(t + i) = *(tl + i);
		*(t + numTriangles) = 0;

		n->setTriangleList(t);
		n->setLeaf(true);

		return;
	}

	n->setLeaf(false);

	// set the split-pos
	n->setSplitPos(split);

	Triangle** left = new (tlmm.request(numLeft)) Triangle*[numLeft];
	int leftCount = 0;
	Triangle** right = new (tlmm.request(numRight)) Triangle*[numRight];
	int rightCount = 0;
	for (int i = 0; i < numTriangles; i++)
	{
		if (tl[i]->getMin(axis) < split)
		{
			left[leftCount] = tl[i];
			leftCount++;
		}
		if (tl[i]->getMax(axis) >= split)
		{
			right[rightCount] = tl[i];
			rightCount++;
		}
	}

	// request a place for the children-nodes
	KdNode* children = new (kdtmm.request(2)) KdNode[2];

	// link the current node to it's children
	n->setLeft(children);

	const AABB leftAABB = aabb.createLeft(split, axis);
	const AABB rightAABB = aabb.createRight(split, axis);
	_buildKdTree(children, left, numLeft, leftAABB, depth + 1);
	_buildKdTree(children + 1, right, numRight, rightAABB, depth + 1);
}

KdNode* buildKdTree(const vector<Triangle*>& triangles, const AABB& aabb)
{
	KdNode* root = new (kdtmm.request(1)) KdNode;

	// create the triangle-list
	const int numTriangles = triangles.size();
	Triangle** tl = new (tlmm.request(numTriangles)) Triangle*[numTriangles];

	// store the triangles in the list
	int index = 0;
	for (vector<Triangle*>::const_iterator it = triangles.begin(); it
			!= triangles.end(); it++)
	{
		tl[index] = *it;
		index++;
	}

	_buildKdTree(root, tl, numTriangles, aabb, 0);
	return root;
}

void printKdTree(const KdNode* n)
{
	std::cout << "Address: " << n << std::endl;
	std::cout << "Is leaf: " << n->isLeaf() << std::endl;
	if (!n->isLeaf())
	{
		std::cout << "Split-position: " << n->getSplitPos() << std::endl;
		std::cout << "Split-axis: " << n->getAxis() << std::endl;
		std::cout << "Address of left child: " << n->getLeft() << std::endl;
		std::cout << "Address of right child: " << n->getRight() << std::endl;
		printKdTree(n->getLeft());
		printKdTree(n->getRight());
	}
	else
	{
		Triangle** cur = n->getTriangleList();
		int num = 0;
		while (*cur != 0)
		{
			num++;
			cur++;
		}
		std::cout << "Address of first triangle: " << n->getTriangleList()
				<< std::endl;
		std::cout << "Number of triangles: " << num << std::endl;
	}
}

static void _drawKdTree(const KdNode* n, int omitAxis, int x, int y,
		const AABB& aabb)
{
	int axis = n->getAxis();

	if (!n->isLeaf())
	{
		float split = n->getSplitPos();

		//		std::cout << split << std::endl;
		//		std::cout << min << std::endl;
		//		std::cout << max << std::endl;
		//		std::cout << std::endl;

		if (omitAxis == X)
		{
			if (axis == Y)
			{
				float min = aabb.getMin().get(Z);
				float max = aabb.getMax().get(Z);
				line(buffer, x + (int)(min * 10), y + (int)(split * 10), x
						+ (int)(max * 10), y + (int)(split * 10), makecol(0,
						255, 0));
			}
			else if (axis == Z)
			{
				float min = aabb.getMin().get(Y);
				float max = aabb.getMax().get(Y);
				line(buffer, x + (int)(split * 10), y + (int)(min * 10), x
						+ (int)(split * 10), y + (int)(max * 10), makecol(0, 0,
						255));
			}
		}
		else if (omitAxis == Y)
		{
			if (axis == X)
			{
				float min = aabb.getMin().get(Z);
				float max = aabb.getMax().get(Z);
				line(buffer, x + (int)(split * 10), y + (int)(min * 10), x
						+ (int)(split * 10), y + (int)(max * 10), makecol(255,
						0, 0));
			}
			else if (axis == Z)
			{
				float min = aabb.getMin().get(X);
				float max = aabb.getMax().get(X);
				line(buffer, x + (int)(min * 10), y + (int)(split * 10), x
						+ (int)(max * 10), y + (int)(split * 10), makecol(0, 0,
						255));
			}
		}
		else if (omitAxis == Z)
		{
			if (axis == X)
			{
				float min = aabb.getMin().get(Y);
				float max = aabb.getMax().get(Y);
				line(buffer, x + (int)(split * 10), y + (int)(min * 10), x
						+ (int)(split * 10), y + (int)(max * 10), makecol(255,
						0, 0));
			}
			else if (axis == Y)
			{
				float min = aabb.getMin().get(X);
				float max = aabb.getMax().get(X);
				line(buffer, x + (int)(min * 10), y + (int)(split * 10), x
						+ (int)(max * 10), y + (int)(split * 10), makecol(0,
						255, 0));
			}
		}

		const AABB leftAABB = aabb.createLeft(split, axis);
		const AABB rightAABB = aabb.createRight(split, axis);
		_drawKdTree(n->getLeft(), omitAxis, x, y, leftAABB);
		_drawKdTree(n->getRight(), omitAxis, x, y, rightAABB);
	}
}

void drawKdTree(const KdNode* n, int omitAxis, int x, int y, const AABB& aabb)
{
	if (omitAxis == X)
		rect(buffer, x + (int)(aabb.getMin().get(Z) * 10), y + (int)(aabb.getMin().get(Y) * 10), x + (int)(aabb.getMax().get(Z) * 10), y + (int)(aabb.getMax().get(Y) * 10), makecol(255, 255, 255));
	else if (omitAxis == Y)
		rect(buffer, x + (int)(aabb.getMin().get(X) * 10), y + (int)(aabb.getMin().get(Z) * 10), x + (int)(aabb.getMax().get(X) * 10), y + (int)(aabb.getMax().get(Z) * 10), makecol(255, 255, 255));
	else if (omitAxis == Z)
		rect(buffer, x + (int)(aabb.getMin().get(X) * 10), y + (int)(aabb.getMin().get(Y) * 10), x + (int)(aabb.getMax().get(X) * 10), y + (int)(aabb.getMax().get(Y) * 10), makecol(255, 255, 255));
	_drawKdTree(n, omitAxis, x, y, aabb);
}

static void intersect(const KdNode* node, Ray& ray)
{
	float distance;
	Triangle** cur = node->getTriangleList();
	while (*cur != 0)
	{
		//			distance = Intersections::rayTriangle(ray.getPoint(),
		//					ray.getVector(), (*cur)->getPoint(0), (*cur)->getPoint(1), (*cur)->getPoint(2));
		distance = Intersections::rayTriangleFast(ray.getPoint(),
				ray.getVector(), (*cur)->triAccel);
		if (distance < ray.getHitDistance())
		{
			ray.setHitDistance(distance);
			ray.setHitTriangle(*cur);
		}
		cur++;
	}
}

static float _recTraverse(KdNode* node, Ray& ray, float t_near, float t_far)
// returns distance to closest hit point
{
	if (node->isLeaf())
	{
		intersect(node, ray);
		return ray.getHitDistance();
	}
	
	// determine the front- and back-child
	KdNode* front;
	KdNode* back;
	if (viewPoint.get(node->getAxis()) > node->getSplitPos())
	{
		front = node->getRight();
		back = node->getLeft();	
	}
	else
	{
		front = node->getLeft();
		back = node->getRight();
	}
	
	float d = (node->getSplitPos() - ray.getPoint().get(node->getAxis())) * ray.getInverseVector().get(node->getAxis());
	d = (d < .0) ? MAXFLOAT : d;

	if (d >= t_far || d < .0)
		return _recTraverse(front, ray, t_near, t_far);
	else if (d <= t_near)
		return _recTraverse(back, ray, t_near, t_far);
	else
	{
		// case three: traverse both sides in turn
		float distance = _recTraverse(front, ray, t_near, d);

		return (distance <= d) ? distance : _recTraverse(back, ray,
				d, t_far);
	}
}

float recTraverse(KdNode* root, Ray& ray, const AABB& aabb)
{
	float t_near = -1000;
	float t_far = 100000.;
	ray.setHitTriangle(0);
	ray.setHitDistance(MAXFLOAT);

	return (Intersections::AABBRay(t_near, t_far, aabb.getMin(), aabb.getMax(),
			ray.getPoint(), ray.getVector())) ? _recTraverse(root, ray, t_near,
			t_far) : MAXFLOAT;
}

/**
 * is the ray intersected?
 */
static bool intersect(const KdNode* node, ShadowRay& shadowRay)
{
	float bestDistance = MAXFLOAT;
	Triangle** cur = node->getTriangleList();
	while (*cur != 0)
	{
		//			distance = Intersections::rayTriangle(ray.getPoint(),
		//					ray.getVector(), (*cur)->getPoint(0), (*cur)->getPoint(1), (*cur)->getPoint(2));
		bestDistance = min(Intersections::rayTriangleFast(shadowRay.getStart(),
				shadowRay.getVector(), (*cur)->triAccel), bestDistance);
		cur++;
	}

	return bestDistance <= shadowRay.getLength();
}

static bool _recTraverseShadow(KdNode* node, ShadowRay& shadowRay, float t_near, float t_far)
// returns distance to closest hit point
{
	if (t_near >= shadowRay.getLength())
		return true;
	
	if (node->isLeaf())
		return !intersect(node, shadowRay);
	
	// determine the front- and back-child
	KdNode* front;
	KdNode* back;
	if (shadowRay.getStart().get(node->getAxis()) > node->getSplitPos())
	{
		front = node->getRight();
		back = node->getLeft();	
	}
	else
	{
		front = node->getLeft();
		back = node->getRight();
	}
	
	float d = (node->getSplitPos() - shadowRay.getStart().get(node->getAxis())) * shadowRay.getInverseVector().get(node->getAxis());
//	d = (d < .0) ? MAXFLOAT : d;
	if (d >= t_far || d < .0)
			return _recTraverseShadow(front, shadowRay, t_near, t_far);
	else if (d <= t_near)
		return _recTraverseShadow(back, shadowRay, t_near, t_far);
	else
	{
		// case three: traverse both sides in turn
		if (!_recTraverseShadow(front, shadowRay, t_near, d))
			return false;

		return _recTraverseShadow(back, shadowRay,
				d, t_far);
	}
}

/**
 * does the shadow-ray reach the target?
 */
bool recTraverseShadow(KdNode* root, ShadowRay& shadowRay, const AABB& aabb)
{
	float t_near = 0;
	float t_far = shadowRay.getLength();

		return _recTraverseShadow(root, shadowRay, t_near,
				t_far);
}

static float _itTraverseShadow(const KdNode* node, ShadowRay& shadowRay, float t_near,
		float t_far)
{		
	while (1)
	{
		while (!node->isLeaf())
		{
			// determine the front- and back-child
			KdNode* front;
			KdNode* back;
			if (shadowRay.getStart().get(node->getAxis()) > node->getSplitPos())
			{
				front = node->getRight();
				back = node->getLeft();	
			}
			else
			{
				front = node->getLeft();
				back = node->getRight();
			}
			
			// traverse ’til next leaf
			float d = (node->getSplitPos() - shadowRay.getStart().get(node->getAxis())) * shadowRay.getInverseVector().get(node->getAxis());
//			d = (d < .0) ? MAXFLOAT : d;
			if (d >= t_far || d < .0)
				node = front;
			else if (d <= t_near)
				node = back;
			else
			{
				// case three: traverse both sides in turn
				stack.push(back, d, t_far);
				node = front;
				t_far = d;
			}
		}

		if (intersect(node, shadowRay))
			return false;
		if (stack.isEmpty())
			return true; // noting else to traverse any more...
		StackEntry s = stack.pop();
		node = s.node;
		t_near = s.near;
		t_far = s.far;
	}
}

bool itTraverseShadow(KdNode* root, ShadowRay& shadowRay, const AABB& aabb)
{
	float t_near = .0;
	float t_far = shadowRay.getLength();
	
	stack.clear();

	return _itTraverseShadow(root, shadowRay, t_near,
					t_far);
}

static float _itTraverse(const KdNode* node, Ray& ray, float t_near,
		float t_far)
{
	while (1)
	{
		while (!node->isLeaf())
		{
			// determine the front- and back-child
			KdNode* front;
			KdNode* back;
			if (viewPoint.get(node->getAxis()) > node->getSplitPos())
			{
				front = node->getRight();
				back = node->getLeft();	
			}
			else
			{
				front = node->getLeft();
				back = node->getRight();
			}
			
			// traverse ’til next leaf
			float d = (node->getSplitPos() - ray.getPoint().get(node->getAxis())) * ray.getInverseVector().get(node->getAxis());
//			d = (d < .0) ? MAXFLOAT : d;
			if (d >= t_far || d < .0)
				node = front;
			else if (d <= t_near)
				node = back;
			else
			{
				// case three: traverse both sides in turn
				stack.push(back, d, t_far);
				node = front;
				t_far = d;
			}
		}
		// have a leaf now
		intersect(node, ray);
		if (ray.getHitTriangle())
			return ray.getHitDistance(); // early ray termination
		if (stack.isEmpty())
			return ray.getHitDistance(); // noting else to traverse any more...
		StackEntry s = stack.pop();
		node = s.node;
		t_near = s.near;
		t_far = s.far;
	}
}

float itTraverse(KdNode* root, Ray& ray, const AABB& aabb)
{
	float t_near = -1000;
	float t_far = 100000.;
	ray.setHitTriangle(0);
	ray.setHitDistance(MAXFLOAT);
	
	stack.clear();

	return (Intersections::AABBRay(t_near, t_far, aabb.getMin(), aabb.getMax(),
				ray.getPoint(), ray.getVector())) ? _itTraverse(root, ray, t_near,
				t_far) : MAXFLOAT;
}
